package org.hepeng.workx.spring.cloud.netflix.zuul.route;

import org.apache.commons.collections.CollectionUtils;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.netflix.zuul.filters.ZuulProperties;
import org.springframework.cloud.netflix.zuul.filters.discovery.DiscoveryClientRouteLocator;
import org.springframework.cloud.netflix.zuul.filters.discovery.ServiceRouteMapper;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;

/**
 * @author he peng
 */
public class CustomSourceRouteLocator extends DiscoveryClientRouteLocator {

    private RouteTableLoader routeTableLoader;

    public CustomSourceRouteLocator(String servletPath, DiscoveryClient discovery,
                                    ZuulProperties properties,
                                    ServiceRouteMapper serviceRouteMapper,
                                    ServiceInstance localServiceInstance,
                                    RouteTableLoader routeTableLoader) {
        super(servletPath, discovery, properties, serviceRouteMapper, localServiceInstance);
        this.routeTableLoader = routeTableLoader;
    }

    @Override
    protected LinkedHashMap<String, ZuulProperties.ZuulRoute> locateRoutes() {
        LinkedHashMap<String, ZuulProperties.ZuulRoute> routesMap = super.locateRoutes();
        Set routeTable = this.routeTableLoader.loadRouteTable();
        if (CollectionUtils.isNotEmpty(routeTable)) {
            Map<String, ZuulProperties.ZuulRoute> routeTableMap =
                    this.routeTableLoader.convert(routeTable);
            routesMap.putAll(routeTableMap);
        }
        return routesMap;
    }
}
